package com.health.es.temp;

import cn.hutool.core.util.ArrayUtil;
import com.health.es.util.ElasticsearchRestTemplate;
import com.health.es.util.SearchHit;
import com.health.es.util.SearchHits;
import com.health.es.util.query.NativeSearchQuery;
import com.health.es.util.query.NativeSearchQueryBuilder;
import com.health.es.util.query.Query;
import org.elasticsearch.index.query.BoolQueryBuilder;
import org.elasticsearch.index.query.QueryBuilder;
import org.elasticsearch.index.query.QueryBuilders;
import org.elasticsearch.search.aggregations.AbstractAggregationBuilder;
import org.elasticsearch.search.aggregations.Aggregation;
import org.elasticsearch.search.aggregations.AggregationBuilders;
import org.elasticsearch.search.aggregations.Aggregations;
import org.elasticsearch.search.aggregations.bucket.MultiBucketsAggregation;
import org.elasticsearch.search.aggregations.bucket.terms.ParsedStringTerms;
import org.elasticsearch.search.aggregations.bucket.terms.TermsAggregationBuilder;
import org.springframework.data.domain.PageRequest;
import org.springframework.data.domain.Pageable;
import org.springframework.data.domain.Sort;

import java.util.Arrays;
import java.util.List;
import java.util.Objects;
import java.util.Set;
import java.util.stream.Collectors;

/**
 * @author nqlz
 * @description: Es 查询辅助工具
 */

public class EsHelper {


    public static ElasticsearchRestTemplate operations;

    /**
     * es 查询 多条记录 selectList
     *
     * @param clazz
     * @param query
     * @param <U>
     * @return
     */
    public static <U> List<U> search(Class<U> clazz, Query query) {
        // spring-data 默认为 unmodifiableList 返回这个list的只读视图lb  add remove 等改变其内容的操作将导致编译不通过。
        return operations.search(query, clazz).map(SearchHit::getContent).toList();
    }

    /**
     * es 查询 单条记录
     *
     * @param clazz
     * @param query
     * @param <U>
     * @return
     */
    public static <U> U searchOne(Class<U> clazz, Query query) {
        SearchHit<U> tSearchHit = operations.searchOne(query, clazz);
        return Objects.isNull(tSearchHit) ? null : tSearchHit.getContent();
    }

    /**
     * 默认查询 10000条数据
     *
     * @param tClass
     * @param builders
     * @param <U>
     * @return
     */
    public static <U> List<U> search(Class<U> tClass, QueryBuilder... builders) {
        QueryBuilder builder = ArrayUtil.length(builders) == 1 ? builders[0] : getBoolQueryBuilder(builders);
        NativeSearchQuery query = new NativeSearchQuery(builder).setPageable(PageRequest.of(0, 10000));
        return search(tClass, query);
    }

    /**
     * sort 排序 默认查询 10000条数据
     *
     * @param tClass
     * @param sort
     * @param builders
     * @param <U>
     * @return
     */
    public static <U> List<U> search(Class<U> tClass, Sort sort, QueryBuilder... builders) {
        NativeSearchQuery query = new NativeSearchQuery(getBoolQueryBuilder(builders))
                .addSort(sort).setPageable(PageRequest.of(0, 10000));
        return search(tClass, query);
    }

    /**
     * selectList wrapper 入参
     *
     * @param tClass
     * @param wrapper
     * @param <U>
     * @return
     */
    public static <U> List<U> search(Class<U> tClass, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        return wrapper.apply(new QueryWrapper<>(tClass)).list();
    }

    /**
     * 查询单条记录
     *
     * @param tClass
     * @param builders
     * @param <U>
     * @return
     */
    public static <U> U searchOne(Class<U> tClass, QueryBuilder... builders) {
        QueryBuilder builder = ArrayUtil.length(builders) == 1 ? builders[0] : getBoolQueryBuilder(builders);
        return searchOne(tClass, new NativeSearchQuery(builder));
    }

    /**
     * 查询当个 wrapper 入参
     *
     * @param tClass
     * @param wrapper
     * @param <U>
     * @return
     */
    public static <U> U searchOne(Class<U> tClass, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        return wrapper.apply(new QueryWrapper<>(tClass)).one();
    }

    /**
     * 分页查询 返回分页对象
     *
     * @param tClass
     * @param builders
     * @param pageable
     * @param <U>
     * @return
     */
    public static <U> EsPage<U> searchPage(Class<U> tClass, Pageable pageable, QueryBuilder... builders) {
        return new EsPage<U>(pageable.getPageNumber(), pageable.getPageSize(), pageable.getSort())
                .getEsPage(tClass, builders);
    }

    /**
     * 使用此分页需配合 注解 Page 使用
     *
     * @param tClass
     * @param sort
     * @param builders
     * @param <U>
     * @return
     */
    public static <U> EsPage<U> searchPage(Class<U> tClass, Sort sort, QueryBuilder... builders) {
        return new EsPage<U>().getEsPage(tClass, sort, builders);
    }

    /**
     * 使用此分页需配合 注解 Page 使用
     *
     * @param tClass
     * @param builders
     * @param <U>
     * @return
     */
    public static <U> EsPage<U> searchPage(Class<U> tClass, QueryBuilder... builders) {
        return new EsPage<U>().getEsPage(tClass, builders);
    }

    /**
     * 分页结果封装 wrapper 入参
     *
     * @param tClass
     * @param wrapper
     * @param <U>
     * @return
     */
    public static <U> EsPage<U> searchPage(Class<U> tClass, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        return wrapper.apply(new QueryWrapper<>(tClass)).listPage();
    }

    /**
     * 获取repository 注意相关类有无相关接口
     *
     * @param tClass
     * @param <U>
     * @return
     */
    public static <U, J> EsCustomRepository<U, J> repository(Class<U> tClass) {
        return EsClassOperations.getRepository(tClass);
    }

    /**
     * 返回原结果
     *
     * @param tClass
     * @param wrapper
     * @param <U>
     * @return
     */
    public static <U> SearchHits<U> prime(Class<U> tClass, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        return wrapper.apply(new QueryWrapper<>(tClass)).prime();
    }

    public static <U> EsPage<U> aggPage(Class<U> tClass, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        return wrapper.apply(new QueryWrapper<>(tClass)).aggPage();
    }

    public static <U> long count(Class<U> tClass, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        return wrapper.apply(new QueryWrapper<>(tClass)).count();
    }

    public static <U> long count(Class<U> tClass, QueryBuilder... builders) {
        QueryBuilder builder = ArrayUtil.length(builders) == 1 ? builders[0] : getBoolQueryBuilder(builders);
        return operations.count(new NativeSearchQuery(builder), tClass);
    }

    /**
     * 快速聚合，不查询字段，不映射结果
     *
     * @param wrapper
     * @param <U>
     * @return
     */
    public static <U, A extends Aggregation> A fastAgg(Class<U> tClass, String aggName, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        QueryWrapper<U> t = wrapper.apply(new QueryWrapper<>(tClass));
        t.nativeSearchQueryBuilder(n -> n.withFields("").withMaxResults(0));
        return Objects.requireNonNull(t.prime().getAggregations()).get(aggName);
    }

    /**
     * 多个 聚合
     *
     * @param tClass
     * @param aggName
     * @param aggregationBuilder
     * @param <U>
     * @param <A>
     * @return
     */
    public static <U, A extends Aggregation> A fastAgg(Class<U> tClass, String aggName, AbstractAggregationBuilder<?>... aggregationBuilder) {
        NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder().withMaxResults(0);
        for (AbstractAggregationBuilder<?> builder : aggregationBuilder) {
            nativeSearchQueryBuilder.addAggregation(builder);
        }
        return operations.search(nativeSearchQueryBuilder.build(), tClass).getAggregations().get(aggName);
    }

    /**
     * 单一字段 聚合查询
     *
     * @param tClass
     * @param fieldName 字段名称
     * @param builders
     * @param <U>
     * @param <A>
     * @return
     */
    public static <U, A extends Aggregation> A fastAgg(Class<U> tClass, String fieldName, QueryBuilder... builders) {
        TermsAggregationBuilder field = AggregationBuilders.terms(fieldName).field(fieldName).size(1000000);
        NativeSearchQueryBuilder nativeSearchQueryBuilder = new NativeSearchQueryBuilder().withMaxResults(0).addAggregation(field);
        nativeSearchQueryBuilder.withQuery(getBoolQueryBuilder(builders));
        Aggregations aggregations = operations.search(nativeSearchQueryBuilder.build(), tClass).getAggregations();
        return aggregations.get(fieldName);
    }

    public static <U, A extends Aggregation> A agg(Class<U> tClass, String aggName, SFunction<QueryWrapper<U>, QueryWrapper<U>> wrapper) {
        QueryWrapper<U> t = wrapper.apply(new QueryWrapper<>(tClass));
        return Objects.requireNonNull(t.prime().getAggregations()).get(aggName);
    }


    /**
     * 获取去重的字段集合
     *
     * @param tClass
     * @param fieldName
     * @param <U>
     * @return
     */
    public static <U> Set<String> getDistinctFields(Class<U> tClass, String fieldName) {
        TermsAggregationBuilder field = AggregationBuilders.terms(fieldName)
                .field(fieldName).size(1000000);
        ParsedStringTerms aggregation = fastAgg(tClass, fieldName, i ->
                i.addAggregation(field).nativeSearchQueryBuilder(n -> n.withMaxResults(0).withCollapseField(fieldName)));
        return aggregation.getBuckets().parallelStream().map(MultiBucketsAggregation.Bucket::getKeyAsString).collect(Collectors.toSet());
    }

    public static <U> Set<String> getDistinctFields(Class<U> tClass, String fieldName, QueryBuilder... builders) {
        TermsAggregationBuilder agg = AggregationBuilders.terms(fieldName)
                .field(fieldName).size(100000);
        QueryBuilder builder = ArrayUtil.length(builders) == 1 ? builders[0] : getBoolQueryBuilder(builders);
        NativeSearchQuery build = new NativeSearchQueryBuilder()
                .withCollapseField(fieldName).withQuery(builder)
                .withMaxResults(0).addAggregation(agg).build();
        ParsedStringTerms aggregation = operations.search(build, tClass).getAggregations().get(fieldName);
        return aggregation.getBuckets().parallelStream().map(MultiBucketsAggregation.Bucket::getKeyAsString).collect(Collectors.toSet());
    }

    public static BoolQueryBuilder getBoolQueryBuilder(QueryBuilder... queryBuilders) {
        return getBoolQueryBuilder(Arrays.asList(queryBuilders));
    }

    public static BoolQueryBuilder getBoolQueryBuilder(List<QueryBuilder> queryBuilders) {
        BoolQueryBuilder boolQueryBuilder = QueryBuilders.boolQuery();
        boolQueryBuilder.must().addAll(queryBuilders);
        return boolQueryBuilder;
    }


}